import os
import cv2
import time
import numpy as np
from YOLO_V4 import darknet


def convertBack(x, y, w, h):
    xmin = int(round(x - (w / 2)))
    xmax = int(round(x + (w / 2)))
    ymin = int(round(y - (h / 2)))
    ymax = int(round(y + (h / 2)))
    return xmin, ymin, xmax, ymax


def cvDrawBoxes(detections, img):
    for detection in detections:
        x, y, w, h = detection[2][0],\
            detection[2][1],\
            detection[2][2],\
            detection[2][3]
        xmin, ymin, xmax, ymax = convertBack(
            float(x), float(y), float(w), float(h))
        pt1 = (xmin, ymin)
        pt2 = (xmax, ymax)
        cv2.rectangle(img, pt1, pt2, (0, 255, 0), 1)
        cv2.putText(img,
                    detection[0].decode() +
                    " [" + str(round(detection[1] * 100, 2)) + "]",
                    (pt1[0], pt1[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    [0, 255, 0], 2)
    return img


netMain = None
metaMain = None
altNames = None


class YOLO(object):
    def __init__(self):
        self.configPath = "/home/vision/Opensource/Broadcast/darknet/VisDrone/vehicle.cfg"
        self.weightPath = "/home/vision/Opensource/Broadcast/darknet/backup/model/vehicle/vehicle_best.weights"
        self.metaPath = "/home/vision/Opensource/Broadcast/darknet/VisDrone/vehicle.data"
        self.netMain = None
        self.metaMain = None
        self.altNames = None

        if not os.path.exists(self.configPath):
            raise ValueError("Invalid config path `" +
                             os.path.abspath(self.configPath)+"`")
        if not os.path.exists(self.weightPath):
            raise ValueError("Invalid weight path `" +
                             os.path.abspath(self.weightPath)+"`")
        if not os.path.exists(self.metaPath):
            raise ValueError("Invalid data file path `" +
                             os.path.abspath(self.metaPath)+"`")
        if self.netMain is None:
            self.netMain = darknet.load_net_custom(self.configPath.encode(
                "ascii"), self.weightPath.encode("ascii"), 0, 1)  # batch size = 1
        if self.metaMain is None:
            self.metaMain = darknet.load_meta(self.metaPath.encode("ascii"))
        if self.altNames is None:
            try:
                with open(self.metaPath) as metaFH:
                    metaContents = metaFH.read()
                    import re
                    match = re.search("names *= *(.*)$", metaContents,
                                      re.IGNORECASE | re.MULTILINE)
                    if match:
                        result = match.group(1)
                    else:
                        result = None
                    try:
                        if os.path.exists(result):
                            with open(result) as namesFH:
                                namesList = namesFH.read().strip().split("\n")
                                self.altNames = [x.strip() for x in namesList]
                    except TypeError:
                        pass
            except Exception:
                pass

    def detectionImg(self,Img):
        Flag = False
        darknet_image = darknet.make_image(darknet.network_width(self.netMain),darknet.network_height(self.netMain), 3)
        m,n,c = Img.shape
        self.srcw = n
        self.srch = m
        w = darknet.network_width(self.netMain)
        h = darknet.network_height(self.netMain)
        Img_rgb = cv2.cvtColor(Img, cv2.COLOR_BGR2RGB)
        frame_resized = cv2.resize(Img_rgb,(w,h),interpolation=cv2.INTER_LINEAR)
        darknet.copy_image_from_bytes(darknet_image, frame_resized.tobytes())
        detections = darknet.detect_image(self.netMain, self.metaMain, darknet_image, thresh=0.25)
        image = cvDrawBoxes(detections, frame_resized)

        image2 = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image3 = cv2.resize(image2, (n, m), interpolation=cv2.INTER_LINEAR)
        for i in range(len(detections)):
            if detections[i][0].decode() == 'vehicle':
                Flag = True
                break
            else:
                Flag = False
        detections = self.DataProcess(detections,image3)
        return image3,Flag,detections

    def DataProcess(self,detections,image3):
        f = []
        w = darknet.network_width(self.netMain)
        h = darknet.network_height(self.netMain)
        src_w =self.srcw
        src_h = self.srch
        coinh = src_h/h
        coinw = src_w/w
        for detect in detections:
            class_name = ['ignored regions','pedestrian','people','bicycle','car','van','truck','tricycle','awning-tricycle','bus','motor','vehicle','others']
            cls = class_name.index(detect[0].decode())
            scort = detect[1]
            x = detect[2][0] * coinw
            y = detect[2][1] * coinh
            w = detect[2][2] * coinw
            h = detect[2][3] * coinh
            xmin, ymin, xmax, ymax = convertBack(
                float(x), float(y), float(w), float(h))
            # cv2.rectangle(image3,(xmin,ymin),(xmax, ymax),(0,0,255),4,cv2.LINE_AA)
            f.append([cls, scort, xmin, ymin, xmax, ymax])
        # cv2.imshow("img", image3)
        # cv2.waitKey(0)
        det = np.array(f)
        return det
    def classname(self):
        class_name = ['ignored regions', 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle',
                      'awning-tricycle', 'bus', 'motor', 'vehicle', 'others']
        return class_name
if __name__ == "__main__":
    img = cv2.imread("/home/vision/Data/DataSet/Poscal_VOC/VOC2007/JPEGImages/000120.JPG")
    app = YOLO()
    result,Flag,det  = app.detectionImg(img)
    cv2.imshow("Yolo", result)
    cv2.waitKey(0)